import open3d as o3d
import numpy as np
import pickle
import sys
import os
# 添加模块路径
sys.path.append(os.path.join(os.path.dirname(__file__), '../SMPL'))
from smpl_np import SMPLModel


#并步抱拳
def set_bingbu_baoquan1_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([0.0, 0.0, 0.0])  # 左髋无旋转
    pose[2] = np.array([0.0, 0.0, 0.0])  # 右髋无旋转
    
    # 膝盖和脚踝保持自然状态，无旋转
    pose[4] = np.array([0.0, 0.0, 0.0])  # 左膝盖
    pose[5] = np.array([0.0, 0.0, 0.0])  # 右膝盖
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为抱拳姿势
    # 左手
    pose[16] = np.array([0.85, 0.0, -1.3])  # 左肩
    pose[18] = np.array([-0.4, -1.5, 0.0])  # 左肘
    
    
    # 右手
    pose[17] = np.array([0.85, -0.0, 1.3])  # 右肩
    pose[19] = np.array([-0.4, 1.5, -0.0])  # 右肘
    
    
    # 脖子
    pose[12] = np.array([0.0, 1.1, 0])  # 头旋转1.5


    return pose


#弓步冲拳
def set_gongbu_chongquan_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([-1.4, 0.0, 0.0])  # 左髋无旋转
    pose[2] = np.array([1.1, 0.0, 0.0])  # 右髋无旋转
    
    # 膝盖和脚踝保持自然状态，无旋转
    pose[4] = np.array([1.4, 0.0, 0.0])  # 左膝盖
    pose[5] = np.array([0.0, 0.0, 0.0])  # 右膝盖
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝
    pose[8] = np.array([-pose[2][0], 0.0, 0.0])  # 右脚踝
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    
    # 左手
    pose[16] = np.array([0.85, 0.0, -1.3])  # 左肩
    pose[18] = np.array([-0.4, -1.5, 0.0])  # 左肘
    
    
    # 右手
    pose[17] = np.array([0.0, 1.4, 0])  # 右肩
    #pose[19] = np.array([-0.4, 1.5, -0.0])  # 右肘
    
    return pose



def set_tanleg_chongquan_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([0.0, 0.0, 0.0])  # 左髋无旋转
    pose[2] = np.array([0.0, 0.0, 0.0])  # 右髋无旋转
    
    # 膝盖调整
    pose[4] = np.array([0.0, 0.0, 0.0])  # 左膝盖无旋转
    pose[5] = np.array([0.0, 0.0, 0.0])  # 右膝盖无旋转
    
    # 脚踝调整
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝无旋转
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝无旋转
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为冲拳姿势
    # 左手（冲拳）
    pose[16] = np.array([0.0, -1.5, 0.0])  # 左肩绕y轴旋转1.0弧度，向前抬起
    
    pose[20] = np.array([0.0, 0.0, 0.0])  # 左手腕无旋转
    pose[22] = np.array([0.0, 0.0, -1.5])  # 左手旋转
    
    # 右手（护胸）
    pose[17] = np.array([0.85, 0.0, 1.3])  # 右肩
    pose[19] = np.array([-0.4, 1.5, 0.0])  # 右肘
    pose[21] = np.array([0.0, 0.0, 1.5])  # 右手
    
    # 腿部调整为弹腿姿势
    # 左腿（支撑）
    pose[1] = np.array([0.0, 0.0, 0.0])  # 左髋
    pose[4] = np.array([0.0, 0.0, 0.0])  # 左膝盖
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝
    
    # 右腿（）
    pose[2] = np.array([-1.5, 0.0, 0.0])  # 右髋
    
    pose[8] = np.array([1.0, 0.0, 0.0])  # 右脚
    
    return pose




def set_mabu_jiada_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([-1.5, 1, 0])  # 左髋无旋转
    pose[2] = np.array([-1.5, -1, 0])  # 右髋无旋转
    
    # 膝盖调整
    pose[4] = np.array([1.5, 0.0, 0.8])  # 左膝盖绕x轴弯曲0.5弧度
    pose[5] = np.array([1.5, 0.0, -0.8])  # 右膝盖绕x轴弯曲0.5弧度
    
    # 脚踝调整
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为架打姿势
    # 左手（架起）
    pose[16] = np.array([0.0, 0.0, 0.8])  # 左肩
    pose[18] = np.array([0.0, 0.0, 1.5])  # 左肘
    pose[20] = np.array([0.0, 0.0, 0.8])  # 左手腕
    
    # 右手（击打）
    pose[17] = np.array([0.0, 0, 0.0])  # 右肩
    pose[19] = np.array([0.0, 0.0, 0])  # 右肘
    pose[23] = np.array([0.0, 0.0, 1.5])  # 右手

    #脖子
    pose[12] = np.array([0.0, -1.0, 0])  # 头旋转1.5
    
    return pose



def set_xiebu_gaida_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([-0.5, 0.0, 0.0])  # 左髋绕x轴向前旋转0.5弧度
    pose[2] = np.array([-1.8, 0.8, 0.0])  # 右髋
    
    # 膝盖调整
    pose[4] = np.array([2.2, 0.0, 0.0])  # 左膝盖绕x轴弯曲0.5弧度
    pose[5] = np.array([1.5, 0.0, 0.0])  # 右膝盖绕x轴弯曲0.5弧度
    
    # 脚踝调整
    pose[7] = np.array([0.2, 0.0, 0.0])  # 左脚踝绕x轴轻微弯曲0.2弧度
    pose[8] = np.array([0.2, 0.0, 0.0])  # 右脚踝绕x轴轻微弯曲0.2弧度
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为盖打姿势
    # 左手（冲拳）
    pose[16] = np.array([0.0, -1.5, 0.0])  # 左肩绕y轴旋转1.0弧度，向前抬起
    
    pose[20] = np.array([0.0, 0.0, 0.0])  # 左手腕无旋转
    pose[22] = np.array([0.0, 0.0, -1.5])  # 左手旋转
    
    # 右手（护胸）
    pose[17] = np.array([0.85, 0.0, 1.3])  # 右肩绕y轴旋转-0.5弧度，向前抬起
    pose[19] = np.array([-0.4, 1.5, 0.0])  # 右肘绕z轴旋转-0.5弧度，弯曲
    pose[21] = np.array([-1.5, 0.0, 00])  # 右手wa无旋转
    pose[23] = np.array([0.0, 0.0, 1.5])  # 右手无旋转
    
    return pose




def set_tixi_chuanzhang_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([-2.2, 0.0, 0.0])  # 左髋绕x轴向前旋转0.5弧度
    pose[2] = np.array([0.0, 0.0, 0.0])  # 右髋无旋转
    
    # 膝盖调整
    pose[4] = np.array([1.5, 0.0, 0.0])  # 左膝盖绕x轴弯曲1.0弧度
    pose[5] = np.array([0.0, 0.0, 0.0])  # 右膝盖无旋转
    
    # 脚踝调整
    pose[7] = np.array([1.2, 0.0, 0.0])  # 左脚踝绕x轴轻微弯曲0.5弧度
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝无旋转
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为穿掌姿势
    # 左手
    pose[16] = np.array([0.0, -0.3, 0.0])  # 左肩无旋转
    pose[18] = np.array([0.0, -2.5, 0.0])  # 左肘无旋转
    pose[20] = np.array([0.0, 0.0, 0.0])  # 左手腕无旋转
    
    # 右手（穿掌）
    pose[17] = np.array([0.0, 0.0, -0.5])  # 右肩
    pose[19] = np.array([-1.5, 0.0, 0])  # 右肘
    pose[21] = np.array([-1.5, 0.0, 0.0])  # 右手腕x旋转
    
    # # 调整左臂为穿掌动作
    # pose[15] = np.array([0.0, 1.0, 0.0])  # 左肩绕y轴旋转1.0弧度，向前抬起
    # pose[16] = np.array([0.0, 0.0, 1.0])  # 左肘绕z轴旋转1.0弧度，伸直
    # pose[18] = np.array([0.0, 0.0, 0.0])  # 左手腕无旋转
    #脖子
    pose[12] = np.array([0.0, -1.0, 0])  # 头旋转1.0
    return pose



def set_pubu_chuanzhang_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    #spine1
    pose[9] = np.array([0.0, 0.0, -0.3])

    # Hips调整
    pose[1] = np.array([-1.5, 1.5, 0.0])  # 左髋无旋转
    pose[2] = np.array([-2, -0.5, 0.0])  # 右髋无旋转
    
    # 膝盖调整
    pose[4] = np.array([0.0, 0.0, 0.0])  # 左膝盖
    pose[5] = np.array([1.8, 0.0, 0.0])  # 右膝盖无旋转
    
    # 脚踝调整
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝无旋转
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为穿掌姿势
    # 左手
    pose[16] = np.array([0.0, 0.0, -0.6])  # 左肩无旋转
    pose[18] = np.array([-1.5, 0.0, 0.0])  # 左肘无旋转
    pose[20] = np.array([-1.5, 0.0, 0.0])  # 左手腕无旋转
    
    # 右手（穿掌）
    pose[17] = np.array([0.0, 0.0, -0.5])  # 右肩
    pose[19] = np.array([-1.5, 0.0, 0])  # 右肘绕
    pose[21] = np.array([-1.5, 0.0, 0.0])  # 右手腕
    #脖子
    pose[12] = np.array([0.0, 1.0, 0])  # 头旋转1.0
    return pose



def set_xubu_tiaozhang_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([-1.6, 0.5, 0.0])  # 左髋无旋转
    pose[2] = np.array([-1.9, -0.3, 0.0])  # 右髋无旋转
    
    # 膝盖调整
    pose[4] = np.array([1.8, 0.0, 0.0])  # 左膝盖无旋转
    pose[5] = np.array([0.8, 0.0, 0.0])  # 右膝盖绕x轴弯曲0.5弧度
    
    # 脚踝调整
    pose[7] = np.array([-0.2, 0.0, 0.0])  # 左脚踝无旋转
    pose[8] = np.array([0.2, 0.0, 0.0])  # 右脚踝绕x轴轻微弯曲0.2弧度
    
    # 躯干保持直立
    #spine1
    pose[3] = np.array([0.0, 1.0, 0.0])
    
    #pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为挑掌姿势
    # 左手
    pose[16] = np.array([0.0, 0.0, 0.3])  # 左肩无旋转
    pose[18] = np.array([0.0, 0.0, 0.0])  # 左肘无旋转
    pose[20] = np.array([0.0, 0.0, -1.3])  # 左手腕无旋转
    
    # 右手（挑掌）
    pose[17] = np.array([0.0, 0.2, 0.0])  # 右肩
    pose[19] = np.array([0.0, 0.0, 0])  # 右肘
    pose[21] = np.array([0.0, 0.0, -1.3])  # 右手腕
    
    #脖子
    pose[12] = np.array([0.0, -1.0, 0])  # 头旋转1.0
    return pose



def set_bingbu_baoquan2_pose():
     # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0.0, 0.0])  # 骨盆保持水平，无旋转
    
    # Hips调整
    pose[1] = np.array([0.0, 0.0, 0.0])  # 左髋无旋转
    pose[2] = np.array([0.0, 0.0, 0.0])  # 右髋无旋转
    
    # 膝盖和脚踝保持自然状态，无旋转
    pose[4] = np.array([0.0, 0.0, 0.0])  # 左膝盖
    pose[5] = np.array([0.0, 0.0, 0.0])  # 右膝盖
    pose[7] = np.array([0.0, 0.0, 0.0])  # 左脚踝
    pose[8] = np.array([0.0, 0.0, 0.0])  # 右脚踝
    
    # 躯干保持直立
    pose[3] = np.array([0.0, 0.0, 0.0])  # 腰部无旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # 胸部无旋转
    
    # 手臂调整为抱拳姿势
    # 左手
    pose[16] = np.array([0.85, 0.0, -1.3])  # 左肩
    pose[18] = np.array([-0.4, -1.5, 0.0])  # 左肘
    
    
    # 右手
    pose[17] = np.array([0.85, -0.0, 1.3])  # 右肩
    pose[19] = np.array([-0.4, 1.5, -0.0])  # 右肘
    
    
    # 脖子
    pose[12] = np.array([0.0, 0, 0])  # 头旋转1.5


    return pose


def set_follow_horse_stance_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0, 0.0])  # 绕x轴旋转0.5弧度
    
    # Hips调整
    pose[1] = np.array([-1.5, 0.0, 0.3])  # 左髋
    pose[2] = np.array([-1.5, 0.0, -0.3])   # 右髋
    
    # spine1和spine2
    pose[3] = np.array([0.0, 0.0, 0.0])  # spine1绕z轴旋转
    pose[6] = np.array([0.0, 0.0, 0.0])  # spine2绕z轴旋转
    
    # knees调整，弯曲膝盖
    pose[5] = np.array([1.57, 0.0, 0.0])  # 左膝绕x轴旋转
    pose[4] = np.array([1.57, 0.0, 0.0])  # 右膝绕x轴旋转
    
    # 肩膀稍微张开，手臂靠腰
    pose[16] = np.array([0.0, 1.4, 0.0])    # 左肩绕y轴旋转
    pose[17] = np.array([0.0, 1.4, 0.0])   # 右肩绕y轴旋转
    
    return pose

#一字马
def set_one_word_horse_pose():
    # 设置姿态参数
    pose = np.zeros((24, 3))
    
    # pelvis
    pose[0] = np.array([0.0, 0, 0.0])  #
    
    # Hips调整
    pose[1] = np.array([1.5, 0.0, 0])  # 左髋绕x轴旋转 x轴旋转1.5
    pose[2] = np.array([-1.5, 0.0, 0])   # 右髋绕x轴旋转-1.5弧度
    
    # l_wrist和r_wrist
    pose[20] = np.array([0.0, 0.0, 1.5])  # l_wrist绕z轴旋转1.5
    pose[21] = np.array([0.0, 0.0, -1.5])  # r_wrist绕z轴旋转-1.5
    
    return pose



def visualize_dung_mabu_posture(model_path, pose_function):
    # 初始化SMPL模型
    model = SMPLModel(model_path)
    
    # 设置姿态参数
    pose = pose_function
    
    
    # 应用姿态参数
    model.set_params(pose=pose)
    
    # 创建点云和网格
    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(model.verts)
    
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(model.verts)
    # 这里假设faces是按照Numpy数组存储的，调整为Open3D所需的格式
     # 检查faces的数据类型
    print("model.faces shape:", model.faces.shape)
    print("model.faces dtype:", model.faces.dtype)
    
    # 将faces转换为int32类型
    # model.faces = model.faces.astype(np.int32)
    # print("After conversion:")
    # print("model.faces shape:", model.faces.shape)
    # print("model.faces dtype:", model.faces.dtype)
    model.faces = model.faces.astype(np.int32)
    mesh.triangles = o3d.utility.Vector3iVector(model.faces)
    mesh.compute_vertex_normals()
    
    # 设置颜色
    point_cloud.paint_uniform_color([0.1, 0.1, 0.7])  # 蓝色点云
    #mesh.paint_uniform_color([0.9, 0.8, 0.7])        # 黄色
    # 高级灰色
    mesh.paint_uniform_color([0.4, 0.5, 0.4])  # 
    
    # 创建关节坐标系网格
    joint_frames = []
    for pos in model.J:
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=pos)
        joint_frames.append(frame)
    
    # 可视化设置
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=800, height=600)
    
    # 添加几何对象
    vis.add_geometry(point_cloud)
    vis.add_geometry(mesh)
    for frame in joint_frames:
        vis.add_geometry(frame)
    
    # 设置视角
    ctr = vis.get_view_control()
    ctr.set_front([0, 0, -1])
    ctr.set_up([0, 1, 0])
    ctr.set_lookat([0, 0, 0])
    ctr.set_zoom(0.5)
    
    # 运行可视化
    vis.run()
    vis.destroy_window()

# 调用函数，传入模型路径
if __name__ == '__main__':
    model_path = '../SMPL/model.pkl'
    #pose = set_bingbu_baoquan1_pose()
    #pose = set_gongbu_chongquan_pose()
    #pose = set_tanleg_chongquan_pose()
    #pose = set_mabu_jiada_pose()
    pose = set_xiebu_gaida_pose()

    #pose = set_tixi_chuanzhang_pose()
    #pose = set_pubu_chuanzhang_pose()

    #pose = set_xubu_tiaozhang_pose()

    #pose = set_bingbu_baoquan2_pose()

    visualize_dung_mabu_posture(model_path, pose)





